"""
特征提取
"""
import cv2
from face_detect import FaceDetect
from facenet import InceptionResnetV1
import torch
import numpy as np
import math
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
class FaceFeat:
    def __init__(self):
        #初始化网络
        face_net = InceptionResnetV1().to(device)
        self.face_net = face_net
        self.name_list = []
        self.db_face_feat = []
        # 加载参数
        face_net.load_state_dict(torch.load('./weights/facenet_best_server.pt',map_location='cpu'))
        #开启验证
        face_net.eval()
        #初始化的时候调用
        self.name_list, self.db_face_feat_list = self.load_db_face_feat()
        pass

    def img_processed(self,face_img):
        #因为(face_img)大小不一致，需要统一设置图像的大小h,w均为112
        if face_img is None:
            print("NO FACE")
        face_img = cv2.resize(face_img, (112, 112))
        # 归一化处理[-1,1]
        face_img = (face_img - 127.5) / 127.5
        #因为opencv打开的图像形状是HWC，而卷积神经网络输入数据的形状是NCHW
        face_img = np.transpose(face_img, (2, 0, 1))
        #(1,3,112,112) --> （N ,C ,H ,W）
        face_img = np.expand_dims(face_img, axis=0)
        return face_img

    def get_face_feat(self,face_img):
        """
        :param face_img:人脸图像
        :return: 特征
        """
        face_img = self.img_processed(face_img)
        #将np数组转换为tensor
        face_img_tensor = torch.Tensor(face_img).to(device)
        #获取特征
        face_feat_tensor = self.face_net(face_img_tensor)
        #将tensor转换为numpy数组
        face_feat = face_feat_tensor.detach().cpu().numpy()
        return face_feat

    def load_db_face_feat(self):
        """
        加载已注册得信息
        :return: name_list,db_face_feat_list
        分析：
        """
        name_list = []
        db_face_feat = []
        f = open('users/user.txt', 'r', encoding='utf-8')
        lines = f.readlines()
        for line in lines:
            feat_str = line.split('|')
            name = feat_str[0]
            face_feat = eval(feat_str[1])
            name_list.append(name)
            db_face_feat.append(face_feat)
        return name_list,db_face_feat
    def cal_similarity(self,face_feat):
        """
        计算相似度
        :param face_feat:被识别人的特征
        :return: 被识别的姓名
        """
        name_list, db_face_feat_list = self.name_list, self.db_face_feat_list
        vector = np.array(list(db_face_feat_list))
        face_feat = np.array(face_feat).flatten().tolist()
        #相似度的计算
        distance = self.diff(face_feat,vector)
        data_list = list(zip(name_list,distance))
        data_list = np.array(data_list,dtype=object)
        #预测出pred_name
        sort_index = np.argsort(data_list[:,1])[::-1]
        pred_name = data_list[sort_index[0]]
        return pred_name

    def diff(self,d1,d2):
        dot = np.sum(np.multiply(d1,d2), axis=1)
        norm = np.linalg.norm(d1,axis=0)*np.linalg.norm(d2,axis=1)
        similarity = dot / norm
        # dist = np.arccos(similarity / math.pi)
        return similarity
# if __name__ == '__main__':
#     face_feat = FaceFeat()

